#!/usr/bin/env python3
# coding=utf-8

import rospy
from std_msgs.msg import String
from collections import defaultdict


class YoloResultAggregator(object):
    """
    订阅YOLO结果与触发指令，按编号(1..6)保留最近一次 (color, state)，
    在收到“return”触发时统计异常颜色并发布 Top-2 颜色。

    - 输入格式："color-number-state"，连字符分隔，大小写不敏感，自动去空白
    - color 仅接受 {red, yellow, orange, blue}，否则忽略整条
    - number 接受数字 1..6 或英文 one..six，超出范围忽略整条
    - state 仅接受 {up, down, normal}，未知则忽略整条
    - 只在收到触发指令(默认话题 /kickball、内容 "return")时发布
      - 若无异常颜色：发布 "None-None"
      - 若仅一种异常颜色：发布 "<color>-None"（颜色为小写）
      - 若两种或更多：按出现次数降序取前两名，次数相同按字母序稳定打破
    - 发布后不清空历史
    - 话题名可通过私有参数配置：
      ~yolo_result_topic, ~kickball_cmd_topic, ~kick_output_topic
    """

    ALLOWED_COLORS = {"red", "yellow", "orange", "blue"}
    ALLOWED_STATES = {"up", "down", "normal"}
    WORD_TO_NUMBER = {
        "one": 1, "two": 2, "three": 3,
        "four": 4, "five": 5, "six": 6,
    }

    def __init__(self):
        rospy.init_node("yolo_result_aggregator_node", anonymous=False)

        # 参数（可重映射）
        self.yolo_result_topic = rospy.get_param("~yolo_result_topic", "/yolo_result")
        self.kickball_cmd_topic = rospy.get_param("~kickball_cmd_topic", "/kickball")
        self.kick_output_topic = rospy.get_param("~kick_output_topic", "/kick")

        # 最近一次 (color, state) 按编号缓存
        self.latest_by_number = {}

        # ROS 通信对象
        self.kick_pub = rospy.Publisher(self.kick_output_topic, String, queue_size=10)
        self.yolo_sub = rospy.Subscriber(self.yolo_result_topic, String, self._on_yolo_result)
        self.cmd_sub = rospy.Subscriber(self.kickball_cmd_topic, String, self._on_kickball_cmd)

        rospy.loginfo("yolo_result_aggregator_node started. yolo_result_topic=%s, kickball_cmd_topic=%s, kick_output_topic=%s",
                      self.yolo_result_topic, self.kickball_cmd_topic, self.kick_output_topic)

    @staticmethod
    def _normalize_token(token):
        return token.strip().lower()

    def _parse_number(self, token):
        t = self._normalize_token(token)
        if t.isdigit():
            try:
                value = int(t)
            except ValueError:
                return None
        else:
            value = self.WORD_TO_NUMBER.get(t)

        if value is None or value < 1 or value > 6:
            return None
        return value

    def _parse_color(self, token):
        t = self._normalize_token(token)
        return t if t in self.ALLOWED_COLORS else None

    def _parse_state(self, token):
        t = self._normalize_token(token)
        return t if t in self.ALLOWED_STATES else None

    def _on_yolo_result(self, msg):
        raw = (msg.data or "").strip()
        if not raw:
            return

        parts = [p.strip() for p in raw.split('-') if p.strip()]
        if len(parts) < 3:
            rospy.logdebug("yolo_result ignored (parts<3): %s", raw)
            return

        color_token, number_token, state_token = parts[0], parts[1], parts[2]

        color = self._parse_color(color_token)
        if color is None:
            rospy.logdebug("ignored: unknown color '%s'", color_token)
            return

        number = self._parse_number(number_token)
        if number is None:
            rospy.logdebug("ignored: invalid number '%s'", number_token)
            return

        state = self._parse_state(state_token)
        if state is None:
            rospy.logdebug("ignored: unknown state '%s'", state_token)
            return

        # 记录最近一次
        self.latest_by_number[number] = (color, state)
        rospy.logdebug("updated number=%s -> (%s, %s)", number, color, state)

    def _on_kickball_cmd(self, msg):
        cmd = (msg.data or "").strip().lower()
        if cmd != "return":
            return

        # 统计异常颜色（state != normal）
        color_count = defaultdict(int)
        for number, (color, state) in self.latest_by_number.items():
            if state != "normal":
                color_count[color] += 1

        if not color_count:
            self._publish_kick_result("None-None")
            return

        # 排序：次数降序，其次颜色字母序升序
        sorted_items = sorted(color_count.items(), key=lambda kv: (-kv[1], kv[0]))

        if len(sorted_items) == 1:
            top_color = sorted_items[0][0]
            self._publish_kick_result(f"{top_color}-None")
            return

        first_color = sorted_items[0][0]
        second_color = sorted_items[1][0]
        self._publish_kick_result(f"{first_color}-{second_color}")

    def _publish_kick_result(self, text):
        msg = String()
        msg.data = text
        self.kick_pub.publish(msg)
        rospy.loginfo("published kick result: %s", text)


def main():
    node = YoloResultAggregator()
    rospy.spin()


if __name__ == "__main__":
    try:
        main()
    except rospy.ROSInterruptException:
        pass


